import json
import base64
from os.path import join, dirname, realpath

from gmutil import SM2PrivateKey, SM2PublicKey
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization

JWK_SM2_SM3 = {
    'kty': 'EC',
    'crv': 'sm2p256v1',
    'x': '',
    'y': '',
    'use': 'sig',
    'alg': 'SGD_SM3_SM2',
    'kid': '01',
}

JWK_RSA_SHA256 = {
    'kty': 'RSA',
    'n': '',
    'e': '',
    'use': 'sig',
    'alg': 'RS256',
    'kid': '02'
}


def create_jwks(uni_name):
    sm2_pri = SM2PrivateKey()
    with open(join(dirname(realpath(__file__)), uni_name + '_sm2_pri.pem'), 'wb') as f:
        f.write(sm2_pri.to_pem())
    sm2_pub = sm2_pri.get_public_key()

    rsa_pri = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    with open(join(dirname(realpath(__file__)), uni_name + '_rsa_pri.pem'), 'wb') as f:
        f.write(rsa_pri.private_bytes(encoding=serialization.Encoding.PEM,
                                      format=serialization.PrivateFormat.PKCS8,
                                      encryption_algorithm=serialization.NoEncryption()))
    rsa_pub = rsa_pri.public_key()

    jwk_01 = dict(JWK_SM2_SM3)
    jwk_01['x'] = base64.urlsafe_b64encode(sm2_pub.point.x_octets).decode()
    jwk_01['y'] = base64.urlsafe_b64encode(sm2_pub.point.y_octets).decode()

    jwk_02 = dict(JWK_RSA_SHA256)
    rsa_pub_nums = rsa_pub.public_numbers()
    jwk_02['n'] = base64.urlsafe_b64encode(rsa_pub_nums.n.to_bytes(length=256, byteorder='big', signed=False)).decode()
    jwk_02['e'] = base64.urlsafe_b64encode(rsa_pub_nums.e.to_bytes(length=3, byteorder='big', signed=False)).decode()

    public_keys = {'keys': [jwk_01, jwk_02]}
    with open(join(dirname(realpath(__file__)), uni_name + '_public_keys.json'), 'w') as f:
        json.dump(public_keys, f, ensure_ascii=False, indent=4)

create_jwks('THU')
create_jwks('UTO')

